# Colin Raffel
# Flux-based noise and tonal representations

import numpy as np
import quantize
import psychoac
import window
#import pylab

# Noise/tonal separator/coder class
class noiseTonalSeparator:

  # Initialization
  def __init__(self, fftSize, Fs, noiseLeak, tonalLeak, tonalSmear, tonalPower ):
    # Set parameters
    self.Fs = Fs
    self.noiseLeak = noiseLeak
    self.tonalLeak = tonalLeak
    self.tonalSmear = tonalSmear
    self.tonalPower = tonalPower
    self.fftSize = fftSize
    
    # Saved variables for measurements, leakiness
    self.Xprevious = np.zeros( fftSize )
    self.Xcurrent = np.zeros( fftSize )
    self.noisePrevious = np.zeros( 25 )
    self.tonalPreviousSpectrum = np.zeros( fftSize )
    self.tonalPrevious = np.zeros( fftSize )
    
    # Make band distribution (used by )
    bandDistribution = psychoac.AssignMDCTLinesFromFreqLimits( self.fftSize, self.Fs )
    # Get various scale factor band measurements
    self.scale = psychoac.ScaleFactorBands( bandDistribution )
    
  # Tick in a frame of samples
  def tick(self, xIn ):
    # Remember current spectrum
    self.Xprevious = self.Xcurrent
    # Get new spectrum
    self.Xcurrent = np.fft.fft( window.SineWindow( xIn ) )
    # Ignore DC
    self.Xcurrent[0] = 0.0
    
  # Tick in a spectrum (for integration with PAC coder)
  def tickSpectrum(self, XIn ):
    # Remember current spectrum
    self.Xprevious = self.Xcurrent
    # Get new spectrum
    self.Xcurrent = XIn
    # Ingore DC
    self.Xcurrent[0] = 0.0
    
  # Returns noise levels for the current input buffer
  def noiseCode( self ):
    # Get magnitude
    XcurrentAbs = np.abs( self.Xcurrent )
    XpreviousAbs = np.abs( self.Xprevious )
    # Find difference
    Xdiff = np.abs( XcurrentAbs - XpreviousAbs )
    # Initialize current noise level array
    noiseCurrent = np.zeros(25)
    # Beginning and end indices
    b = self.scale.lowerLine
    e = self.scale.upperLine
    # Get output array
    for i in np.arange(0, 25):
      # Get noise level for this band - NOTE - 128 is a magic value, should be calculated
      noiseCurrent[i] = np.sqrt(np.dot(Xdiff[b[i]:e[i]], Xdiff[b[i]:e[i]])/self.scale.nLines[i])/128.0
    # Leaky integrator
    noiseLevels = self.noiseLeak*noiseCurrent + (1-self.noiseLeak)*self.noisePrevious
    # Save for next time
    self.noisePrevious = noiseCurrent
    return noiseLevels
  
  # Return a frame of the tonal component
  def getTonal( self ):
    # Get magnitude of smeared spectra
    XcurrentSmeared = np.abs( self.tonalSmear*self.Xcurrent + .5*(1-self.tonalSmear)*np.roll(self.Xcurrent, 1) + .5*(1-self.tonalSmear)*np.roll(self.Xcurrent, -1) )
    XpreviousSmeared = np.abs( self.tonalSmear*self.Xprevious + .5*(1-self.tonalSmear)*np.roll(self.Xprevious, 1) + .5*(1-self.tonalSmear)*np.roll(self.Xprevious, -1) )
    # Ignore DC
    XcurrentSmeared[0] = 0.0
    XpreviousSmeared[0] = 0.0
    # Get maximum values
    Xmax = np.greater_equal(XcurrentSmeared, XpreviousSmeared)*XcurrentSmeared + np.greater(XpreviousSmeared, XcurrentSmeared)*XpreviousSmeared
    # Get minimum values
    Xmin = np.less(XcurrentSmeared, XpreviousSmeared)*XcurrentSmeared + np.less_equal(XpreviousSmeared, XcurrentSmeared)*XpreviousSmeared
    # Mask X1 with power'd ratio
    tonalCurrent = self.Xcurrent*np.power( np.nan_to_num(Xmin/Xmax), self.tonalPower )
    # Return leaky ifft of current and previous spectrum
    tonalOutput = np.fft.ifft( self.tonalLeak*tonalCurrent + (1-self.tonalLeak)*self.tonalPreviousSpectrum )
    # A kludge: cyclical overlap-add, assuming a periodic-ish signal!
    tonalFinalOutput = np.append( self.tonalPrevious[self.fftSize/2.0:self.fftSize], np.zeros(self.fftSize/2.0) ) + tonalOutput + np.append( np.zeros(self.fftSize/2.0), tonalOutput[0:self.fftSize/2.0] )
    # Save for next time
    self.tonalPreviousSpectrum = tonalCurrent
    self.tonalPrevious = tonalOutput
    return tonalFinalOutput
    
# Noise level decoder/synthesizer class
class noiseDecoder:

  # Initialization
  def __init__(self, fftSize, Fs, noiseSmear):
    # Set paramters
    self.Fs = Fs
    # Doubled for overlap/add
    self.fftSize = 2*fftSize
    self.noiseSmear = noiseSmear
    
    # For overlap/add
    self.previousOutput = np.zeros( 2*fftSize )
    
    # Make new band distribution
    bandDistribution = psychoac.AssignMDCTLinesFromFreqLimits( 2*fftSize, Fs )
    # Get various scale factor band measurements
    self.scale = psychoac.ScaleFactorBands(bandDistribution)
    
  def noiseDecode( self, noiseLevels ):
    # Generate random reals
    randReals = 2*np.random.rand(self.fftSize/2.0) - 1
    randImag = 2*np.random.rand(self.fftSize/2.0) - 1
    # Scale
    randReals = (self.fftSize/2.0)*randReals/np.sqrt(2.0)
    randImag = (self.fftSize/2.0)*randImag/np.sqrt(2.0)
    # Get interpolated level values
    interpLevels = np.interp( np.arange(self.fftSize/2.0), np.floor( (self.scale.lowerLine + self.scale.upperLine)/2.0 ), noiseLevels )
    # Scale randoms by it!
    scaledSpectrum = interpLevels*(randReals + 1j * randImag)
    # Perform ifft
    currentOutput = np.fft.ifft(np.append(scaledSpectrum, np.flipud(scaledSpectrum)))
    # Window it
    currentOutput = window.SineWindow(currentOutput)
    # Overlap-add
    output = currentOutput[0:self.fftSize/2] + self.previousOutput[self.fftSize/2:self.fftSize]
    # Save for next time
    self.previousOutput = currentOutput
    return np.real(output)

################################################
################################################
# CODE BELOW IS DEPRECATED, FOR REFERENCE ONLY #
################################################
################################################


# Code x0, current input, and x1's, previous input, noise difference
def noiseCode( x0, x1, zpf = 8, Fs = 44100 ):
  # Zero-pad dramatically
  x0 = np.append( x0, np.zeros(np.size(x0)*(zpf-1)) )
  x1 = np.append( x1, np.zeros(np.size(x1)*(zpf-1)) )
  # Get FFTs
  X1 = np.fft.fft(x1)
  X0 = np.fft.fft(x0)
  # Ignore DC
  X1[0] = 0.0
  X0[0] = 0.0
  # Ensure we have magnitude
  X1 = np.abs(X1)
  X0 = np.abs(X0)
  # Find difference
  Xdiff = np.abs(X1 - X0)
  # Make new band distribution
  bandDistribution = psychoac.AssignMDCTLinesFromFreqLimits( np.size(X1), Fs )
  # Get various scale factor band measurements
  scale = psychoac.ScaleFactorBands(bandDistribution)
  noiseLevels = np.zeros(25)
  # Get output array
  for i in np.arange(0, 25):
    # Beginning and end indices
    b = scale.lowerLine[i]
    e = scale.upperLine[i]
    # Get RMS of difference for this band
    RMSdiff = np.sqrt(np.dot(Xdiff[b:e], Xdiff[b:e])/scale.nLines[i])
    # Get noise level for this band
    noiseLevels[i] = RMSdiff/128.0
#  print np.sqrt(np.dot(Xdiff, Xdiff)/scale.nLines[i])/np.sqrt(np.dot(X0, X0)/scale.nLines[i])
#  print noiseLevels
  return noiseLevels
  
# Code x0, current input, and x1's, previous input, noise difference
def noiseCodeLeaky( x0, x1, previousNoiseLevels, g, zpf = 8, Fs = 44100 ):
  # Do a standard noise code
  currentNoiseLevels = noiseCode( x0, x1, zpf, Fs )
  # Leaky integrator
  noiseLevels = g*currentNoiseLevels + (1-g)*previousNoiseLevels
  return noiseLevels
  
def noiseDecode( noiseLevels, fftSize, smear = 0.0, Fs = 44100 ):
  # Generate random reals
  randReals = 2*np.random.rand(fftSize/2.0) - 1
  randImag = 2*np.random.rand(fftSize/2.0) - 1
  # Scale
  randReals = (fftSize/2.0)*randReals/np.sqrt(2.0)
  randImag = (fftSize/2.0)*randImag/np.sqrt(2.0)
  # Make new band distribution
  bandDistribution = psychoac.AssignMDCTLinesFromFreqLimits( fftSize, Fs )
  # Get various scale factor band measurements
  scale = psychoac.ScaleFactorBands(bandDistribution)
  scaledSpectrum = np.zeros(fftSize/2.0) + 1j * np.zeros(fftSize/2.0)
  # Scale according to noise ratio value
  for i in np.arange(0, 25):
    # Beginning and end indices
    b = scale.lowerLine[i]
    e = scale.upperLine[i]
    # Calculate smeared value
    smearedValue = .5*smear*np.greater(i-1, 0)*noiseLevels[i-1] + (1 - smear)*noiseLevels[i] + .5*smear*np.less(i+1, 25)*noiseLevels[(i+1)%24]
    # Fill in scaled values with smearing if needed
    scaledSpectrum[b:e] = smearedValue*(randReals[b:e] + 1j * randImag[b:e])
  # Perform ifft
  output = np.fft.ifft(np.append(scaledSpectrum, np.flipud(scaledSpectrum)))
#  print max(np.real(output))
  return np.real(output)
  
def noiseDecodeWindowed( noiseLevels, fftSize, smear = 0.0, Fs = 44100 ):
  # Generate scaled noise
  scaledNoise = noiseDecode( noiseLevels, fftSize*2, smear, Fs )
  # Window it
  return window.SineWindow(scaledNoise)

# Get the tonal component
def getTonal( x0, x1, threshold = .9, smear = 0.5, power = 100.0):
  # Get FFTs
  X1 = np.fft.fft(x1)
  X0 = np.fft.fft(x0)
  # Ensure we have magnitude
  X1mag = np.abs(X1)
  X0mag = np.abs(X0)
  # Ignore DC
  X1mag[0] = 0.0
  X0mag[0] = 0.0
  # Smear
  X1mag = smear*X1mag + .5*(1-smear)*np.roll(X1mag, 1) + .5*(1-smear)*np.roll(X1mag, -1)
  X0mag = smear*X0mag + .5*(1-smear)*np.roll(X0mag, 1) + .5*(1-smear)*np.roll(X0mag, -1)
  # Ignore DC (again)
  X1mag[0] = 0.0
  X0mag[0] = 0.0
  # Get maximum values
  Xmax = np.greater_equal(X1mag, X0mag)*X1mag + np.greater(X0mag, X1mag)*X0mag
  
#  pylab.plot(Xmax, 'r.')
#  pylab.hold(True)
#  pylab.plot(np.greater_equal(X1mag, X0mag)*X1mag, 'g:')
#  pylab.plot(np.greater(X0mag, X1mag)*X0mag, 'b--')
#  pylab.plot(np.zeros(np.size(X0)), 'k')
#  pylab.hold(False)
#  pylab.show()

  # Get minimum values
  Xmin = np.less(X1mag, X0mag)*X1mag + np.less_equal(X0mag, X1mag)*X0mag
  
  # Deprecated code
  # Get ratios
  # Xratio = X1mag/X0mag
  # Check if the ratio is in range for a stationary component
  # XinRange = np.logical_and( np.greater(Xratio, threshold), np.less(Xratio, 1/threshold) )
  # Mask X1
  # X1 = X1*XinRange
  
  # Mask X1 with power'd ratio
  X1 = X1*np.power( np.nan_to_num(Xmin/Xmax), power )
  # Take IFFT
  output = np.fft.ifft(X1)
  return output